

from config.model_config import get_base_open_ai
from config.embedding_config import bge_model_name
import numpy as np
from numpy import dot
from numpy.linalg import norm


client = get_base_open_ai()


# 余弦距离, 值越大越相似
def cos_sim(a, b):
    return dot(a, b) / (norm(a) * norm(b))

# 欧式距离 越小越相似
def l2(a, b):
    x = np.asarray(a) - np.asarray(b)
    return norm(x)

def get_embedding_data(texts):
    data = client.embeddings.create(
        model=bge_model_name, #填写需要调用的模型编码
        input=texts,
    ).data
    return [x.embedding for x in data]


query = '美食'
document_texts = [
            "美食非常美味，服务员也很友好。",
            "这部电影既刺激又令人兴奋。",
            "阅读书籍是扩展知识的好方法。"
        ]

query_vec = get_embedding_data([query])[0]
doc_vec = get_embedding_data(document_texts)

print('余弦距离:')
for vec in doc_vec:
    print(cos_sim(query_vec, vec))


print('欧式距离:')
for vec in doc_vec:
    print(l2(query_vec, vec))

